/*
 * Copyright 2019 The FATE Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.webank.ai.fate.serving.bean;

import com.google.common.collect.Lists;
import com.google.protobuf.ByteString;
import com.webank.ai.fate.api.mlmodel.manager.ModelServiceProto;
import com.webank.ai.fate.api.serving.InferenceServiceProto;
import com.webank.ai.fate.serving.core.bean.BatchInferenceRequest;
import com.webank.ai.fate.serving.core.bean.InferenceRequest;
import com.webank.ai.fate.serving.core.utils.JsonUtil;
import org.apache.commons.lang3.StringUtils;
import org.junit.BeforeClass;
import org.junit.FixMethodOrder;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.junit.runners.MethodSorters;

import java.io.UnsupportedEncodingException;
import java.net.URL;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@RunWith(JUnit4.class)
@FixMethodOrder(MethodSorters.NAME_ASCENDING)
public class Testing {

    InferenceClient inferenceClient = new InferenceClient("localhost", 8000);

    @BeforeClass
    public static void init() {

    }

    @Test
    public void test() {
        String str = " assd ";
        System.out.println(StringUtils.trim(str));
    }

    @Test
    public void testLoadCacheModel() {
        /*test_model_load("model_2020071410461638392952_guest#9999#arbiter-10000#guest-9999#host-10000#model_cache", "guest", "2020071410461638392952");
        test_model_load("model_2020071410461638392952_host#10000#arbiter-10000#guest-9999#host-10000#model_cache", "host", "2020071410461638392952");
        test_model_Bind("lr-tag-value", "2020071410461638392952");*/

        test_model_load("model_2020071410544991772353_guest#9999#arbiter-10000#guest-9999#host-10000#model_cache", "guest", "2020071410544991772353");
        test_model_load("model_2020071410544991772353_host#10000#arbiter-10000#guest-9999#host-10000#model_cache", "host", "2020071410544991772353");
        test_model_Bind("sbt-tag-value", "2020071410544991772353");
    }

    @Test
    public void testLRModel() {
        test_model_load("model_2020071510534424818871_guest#9999#arbiter-10000#guest-9999#host-10000#model_cache", "guest", "2020071510534424818871");
        test_model_load("model_2020071510534424818871_host#10000#arbiter-10000#guest-9999#host-10000#model_cache", "host", "2020071510534424818871");
        test_model_Bind("lr0715", "2020071510534424818871");
    }

    public void test_model_load(String filename, String currentRole, String tablename) {
        URL resource = Testing.class.getClassLoader().getResource(filename);
        String filePath = resource.getPath().replaceAll("%23", "#");
        ModelServiceProto.PublishRequest.Builder publicRequestBuilder = ModelServiceProto.PublishRequest.newBuilder();
        ModelServiceProto.PublishRequest publishRequest = publicRequestBuilder.setLocal(ModelServiceProto.LocalInfo.newBuilder()
                .setRole(currentRole).setPartyId(currentRole.equalsIgnoreCase("guest") ? "9999" : "10000").build())
                .putRole("guest", ModelServiceProto.Party.newBuilder().addPartyId("9999").build())
                .putRole("arbiter", ModelServiceProto.Party.newBuilder().addPartyId("10000").build())
                .putRole("host", ModelServiceProto.Party.newBuilder().addPartyId("10000").build())
                .putModel("guest", ModelServiceProto.RoleModelInfo.newBuilder().putRoleModelInfo("9999",
                        ModelServiceProto.ModelInfo.newBuilder().setTableName(tablename).setNamespace("guest#9999#arbiter-10000#guest-9999#host-10000#model").build()).build())
                .putModel("host", ModelServiceProto.RoleModelInfo.newBuilder().putRoleModelInfo("10000",
                        ModelServiceProto.ModelInfo.newBuilder().setTableName(tablename).setNamespace("host#10000#arbiter-10000#guest-9999#host-10000#model").build()).build())
                .putModel("arbiter", ModelServiceProto.RoleModelInfo.newBuilder().putRoleModelInfo("10000",
                        ModelServiceProto.ModelInfo.newBuilder().setTableName(tablename).setNamespace("arbiter#10000#arbiter-10000#guest-9999#host-10000#model").build()).build())
                .setLoadType("CACHE")
                .setFilePath(filePath)
                .build();

        inferenceClient.load(publishRequest);
    }

    public void test_model_Bind(String serviceId, String tableName) {
        ModelServiceProto.PublishRequest.Builder publicRequestBuilder = ModelServiceProto.PublishRequest.newBuilder();
        ModelServiceProto.PublishRequest publishRequest = publicRequestBuilder.setLocal(ModelServiceProto.LocalInfo.newBuilder().setRole("guest").setPartyId("9999").build())
                .putRole("guest", ModelServiceProto.Party.newBuilder().addPartyId("9999").build())
                .putRole("arbiter", ModelServiceProto.Party.newBuilder().addPartyId("10000").build())
                .putRole("host", ModelServiceProto.Party.newBuilder().addPartyId("10000").build())
                .putModel("guest", ModelServiceProto.RoleModelInfo.newBuilder().putRoleModelInfo("9999",
                        ModelServiceProto.ModelInfo.newBuilder().setTableName(tableName).setNamespace("guest#9999#arbiter-10000#guest-9999#host-10000#model").build()).build())
                .putModel("host", ModelServiceProto.RoleModelInfo.newBuilder().putRoleModelInfo("10000",
                        ModelServiceProto.ModelInfo.newBuilder().setTableName(tableName).setNamespace("host#10000#arbiter-10000#guest-9999#host-10000#model").build()).build())
                .putModel("arbiter", ModelServiceProto.RoleModelInfo.newBuilder().putRoleModelInfo("10000",
                        ModelServiceProto.ModelInfo.newBuilder().setTableName(tableName).setNamespace("arbiter#10000#arbiter-10000#guest-9999#host-10000#model").build()).build())
                .setLoadType("FILE")
                .setServiceId(serviceId)
                .build();

        inferenceClient.bind(publishRequest);
    }

    private Map<String, Object> getFeatureData() {
        Object[] data = {-0.161357, 0.822813, -0.031609, -0.248363, 1.662757, 1.81831, 1.280035, 1.391616, 2.389857, 1.28865};
        Map<String, Object> featureDataMap = new HashMap<>();
        for (int i = 0; i < 10; i++) {
            featureDataMap.put("x" + i, data[i]);
        }
        return featureDataMap;
    }

    @Test
    public void test_Inference() throws UnsupportedEncodingException {
        InferenceRequest inferenceRequest = new InferenceRequest();
        inferenceRequest.setServiceId("sbt-tag-value");
        inferenceRequest.getFeatureData().putAll(getFeatureData());
        inferenceRequest.getSendToRemoteFeatureData().put("device_id", "8_tv");
        inferenceRequest.getSendToRemoteFeatureData().put("phone_num", "122222222");

//        while (true) {
            InferenceServiceProto.InferenceMessage.Builder inferenceMessageBuilder = InferenceServiceProto.InferenceMessage.newBuilder();
            inferenceMessageBuilder.setBody(ByteString.copyFrom(JsonUtil.object2Json(inferenceRequest), "UTF-8"));
            InferenceServiceProto.InferenceMessage inferenceMessage = inferenceMessageBuilder.build();
            InferenceServiceProto.InferenceMessage resultMessage = inferenceClient.inference(inferenceMessage);
            System.err.println("result ==================" + new String(resultMessage.getBody().toByteArray()));
//        }
    }

    @Test
    public void test_04_BatchInference() {
//        while (true) {
            try {
                BatchInferenceRequest batchInferenceRequest = new BatchInferenceRequest();
                batchInferenceRequest.setCaseId(Long.toString(System.currentTimeMillis()));
                List<BatchInferenceRequest.SingleInferenceData> singleInferenceDataList = Lists.newArrayList();
                for (int i = 0; i < 5; i++) {
                    BatchInferenceRequest.SingleInferenceData singleInferenceData = new BatchInferenceRequest.SingleInferenceData();

                    singleInferenceData.getFeatureData().putAll(getFeatureData());

                    singleInferenceData.getSendToRemoteFeatureData().put("device_id", "8_tv");
                    singleInferenceData.getSendToRemoteFeatureData().put("phone_num", "122222222");

                    singleInferenceData.setIndex(i);
                    singleInferenceDataList.add(singleInferenceData);

                }
                batchInferenceRequest.setBatchDataList(singleInferenceDataList);
                batchInferenceRequest.setServiceId("sbt-tag-value");
                InferenceServiceProto.InferenceMessage.Builder inferenceMessageBuilder =
                        InferenceServiceProto.InferenceMessage.newBuilder();
                String contentString = JsonUtil.object2Json(batchInferenceRequest);
//                System.err.println("send data ===" + contentString);
                try {
                    inferenceMessageBuilder.setBody(ByteString.copyFrom(contentString, "UTF-8"));

                } catch (UnsupportedEncodingException e) {
                    e.printStackTrace();
                }

                InferenceServiceProto.InferenceMessage inferenceMessage = inferenceMessageBuilder.build();

                System.err.println(inferenceMessage.getBody());

                InferenceServiceProto.InferenceMessage resultMessage = inferenceClient.batchInference(inferenceMessage);

                System.err.println("result ==================" + new String(resultMessage.getBody().toByteArray()));
                Thread.sleep(300);
            } catch (Exception e) {

            }
//        }
    }

}